# coding:utf-8
import math
import numpy as np
from PIL import Image
from keras import Sequential
from keras.datasets import mnist
from keras.layers import Activation, Dense, BatchNormalization, Reshape, UpSampling2D, Conv2D, MaxPooling2D, Flatten

#  https://www.jianshu.com/p/f1485a7deb47
# 定义生成模型
from keras.optimizers import SGD


def generator_model():
    model = Sequential()
    model.add(Dense(1024,input_dim=100))
    model.add(Activation('tanh'))

    model.add(Dense(128*7*7))
    # 添加一个批量归一化层，该层在每个batch上将前一层的激活值重新规范化，即使得其输出数据的均值接近0，其标准差接近1
    model.add(BatchNormalization())
    model.add(Activation('tanh'))

    # Reshape层用来将输入shape转换为特定的shape，将含有128*7*7个元素的向量转化为7×7×128张量
    model.add(Reshape((7,7,128),input_shape=(128*7*7,)))
    # 2维上采样层，即将数据的行和列分别重复2次
    model.add(UpSampling2D(size=(2,2)))

    model.add(Conv2D(64,(5,5),padding='same',activation='tanh'))
    model.add(UpSampling2D(size=(2,2)))

    model.add(Conv2D(1,(5,5),padding='same',activation='tanh'))

    return model


# 定义个判别模型
def discriminator_model():
    model = Sequential()

    model.add(Conv2D(64,(5,5),input_shape=(28,28,1)))
    model.add(Activation('tanh'))
    # 为空域信号施加最大值池化，pool_size取（2，2）代表使图片在两个维度上均变为原长的一半
    model.add(MaxPooling2D(pool_size=(2,2)))

    model.add(Conv2D(128,(5,5),activation='tanh'))
    model.add(MaxPooling2D(pool_size=(2,2)))
    # #Flatten层把多维输入一维化，常用在从卷积层到全连接层的过渡
    model.add(Flatten())
    model.add(Dense(1024))
    model.add(Activation('tanh'))

    model.add(Dense(1))
    model.add(Activation("sigmoid"))
    return model


# 拼接两个类别模型
def generator_containing_discriminator(g, d):
    """
     #将前面定义的生成器架构和判别器架构组拼接成一个大的神经网络，用于判别生成的图片
    """
    model = Sequential()
    model.add(g)
    # FAQ： 这里应该先训练判别器，再训练生成器
    # 两个模型 需要固定一个再训练一个
    # 先添加生成器架构，再令d不可训练，即固定d
    # 因此在给定d的情况下训练生成器，即通过将生成的结果投入到判别器进行辨别而优化生成器

    d.trainable=False
    model.add(d)
    return model


# 生成拼接的图片（即将一个batch所有生成图片放到一个图片中）：
def combine_images(generated_images):
    num = generated_images.shape[0]
    width = int(math.sqrt(num))
    height = int(math.ceil(float(num)/width))
    shape = generated_images.shape[1:3]
    image = np.zeros((height*shape[0],
                      width*shape[1]),
                     dtype=generated_images.dtype)
    for index, img in enumerate(generated_images):
        i = int(index/width)
        j = index % width
        image[i*shape[0]:(i+1)*shape[0],j*shape[1]:(j+1)*shape[1]] = \
                                                 img[:,:,0]
    return image


# 训练
def train(BATCH_SIZE):
    # 下载的地址为：https://s3.amazonaws.com/img-datasets/mnist.npz
    (X_train, y_train), (X_test, y_test) = mnist.load_data()
    # print(type(X_train),X_train.shape,X_test.shape)
    # print(type(y_train),y_train.shape)
    # print(X_train[0])
    X_train = (X_train.astype(np.float32) - 127.5) / 127.5
    X_train = X_train[:, :, :, None]  # None将3维的X_train扩展为4维
    X_test = X_test[:, :, :, None]

    d = discriminator_model()
    g = generator_model()
    d_on_g = generator_containing_discriminator(g,d)
    # 定义生成器模型判别器模型更新所使用的优化算法及超参数
    d_optim = SGD(lr=0.001, momentum=0.9, nesterov=True)
    g_optim = SGD(lr=0.001, momentum=0.9, nesterov=True)

    # 编译三个神经网络并设置损失函数和优化算法，其中损失函数都是用的是二元分类交叉熵函数。编译是用来配置模型学习过程的
    g.compile(loss='binary_crossentropy',optimizer='SGD')
    d_on_g.compile(loss='binary_crossentropy',optimizer=g_optim)
    # FAQ: 生成器g 没有训练？
    # 前一个架构在固定判别器的情况下训练了生成器，所以在训练判别器之前先要设定其为可训练。
    d.trainable = True
    d.compile(loss='binary_crossentropy', optimizer=d_optim)

    for epoch in range(30):
        print("epoch is",epoch)
        # 计算一个epoch所需要的迭代数量，即训练样本数除批量大小数的值取整；其中shape[0]就是读取矩阵第一维度的长度
        print("Number of batches", int(X_train.shape[0] / BATCH_SIZE))

        # 在一个epoch内进行迭代训练
        for index in range(int(X_train.shape[0] / BATCH_SIZE)):
            # 随机生成的噪声服从均匀分布，且采样下界为-1、采样上界为1，输出BATCH_SIZE×100个样本；即抽取一个批量的随机样本
            noise = np.random.uniform(-1, 1, size=(BATCH_SIZE, 100))
            # 抽取一个批量的真实图片
            image_batch = X_train[index*BATCH_SIZE:(index+1)*BATCH_SIZE]

            # 生成的图片使用生成器对随机噪声进行推断；verbose为日志显示，0为不在标准输出流输出日志信息，1为输出进度条记录
            generated_images = g.predict(noise,verbose=0)
            print(np.shape(generated_images))

            # 每经过100次迭代输出一张生成的图片
            if index % 100 == 0:
                image = combine_images(generated_images)
                image = image*127.5+127.5
                Image.fromarray(image.astype(np.uint8)).save(
                    "./GAN"+str(epoch)+"_"+str(index)+".png")
            # 将真实的图片和生成的图片以多维数组的形式拼接在一起，真实图片在上，生成图片在下
            X = np.concatenate((image_batch,generated_images))
            # FIXME:这里用平滑标签0.7替换1
            # 生成图片真假标签，即一个包含两倍批量大小的列表；前一个批量大小都是1，代表真实图片，后一个批量大小都是0，代表伪造图片
            y = [1] * BATCH_SIZE + [0] * BATCH_SIZE

            # 判别器的损失；在一个batch的数据上进行一次参数更新
            d_loss = d.train_on_batch(X,y)
            print("batch %d d_loss : %f" % (index, d_loss)) # 理论上，d的loss越来越大，因为生成图片和真实图片越来越像

            # 随机生成的噪声服从均匀分布
            noise = np.random.uniform(-1, 1, (BATCH_SIZE, 100))

            d.trainable = False

            # 计算生成器损失；在一个batch的数据上进行一次参数更新
            # 生成器的目标是愚弄辨别器蒙混过关，需要达到的目标是对于生成的图片，输出为1(正好和鉴别器相反).
            g_loss = d_on_g.train_on_batch(noise,[1]*BATCH_SIZE)

            d.trainable = True
            print("batch %d g_loss : %f" % (index, g_loss)) # 理论上，g的loss越来越小，因为生成图像越接近真实，生成图像的label接近1

            if index % 100 == 0:
                g.save_weights('generator',True)
                d.save_weights('discriminator',True)


def generate(BATCH_SIZE,nice=False):
    g = generator_model()
    g.compile(loss='binary_crossentropy', optimizer="SGD")
    g.load_weights('generator')
    if nice:
        d = discriminator_model()
        d.compile(loss='binary_crossentropy', optimizer="SGD")
        d.load_weights('discriminator')
        noise = np.random.uniform(-1, 1, (BATCH_SIZE * 20, 100))
        generated_images = g.predict(noise, verbose=1)
        d_pret = d.predict(generated_images,verbose=1)
        index = np.arange(0,BATCH_SIZE*20)
        index.resize((BATCH_SIZE*20,1))
        pre_with_index = list(np.append(d_pret, index, axis=1))
        pre_with_index.sort(key=lambda x: x[0], reverse=True)
        nice_images = np.zeros((BATCH_SIZE,)+generated_images.shape[1:3],dtype=np.float32)
        nice_images = nice_images[:,:,:,None]
        for i in range(BATCH_SIZE):
            idx = int(pre_with_index[i][1])
            nice_images[i, :, :, 0] = generated_images[idx, :, :, 0]
        image = combine_images(nice_images)
    else:
        noise = np.random.uniform(-1, 1, (BATCH_SIZE, 100))
        generated_images = g.predict(noise, verbose=0)
        image = combine_images(generated_images)
    image = image * 127.5 + 127.5
    Image.fromarray(image.astype(np.uint8)).save(
        "./GAN/generated_image.png")
train(100)

